import yaml
import numpy as np
import cv2


class location():
    def __init__(self):
        self.calibration_data = self.read_yaml('calibration_matrix.yaml')
        #if self.calibration_data is None:
        #    print("无法加载内参矩阵和畸变系数，程序终止。")
        #    return
        self.camera_matrix = self.check_and_reshape(self.calibration_data['camera_matrix'], (3, 3), "相机内参矩阵")
        self.dist_coeff = self.check_and_reshape(self.calibration_data['dist_coeff'], (5,), "畸变系数")
        self.dist_coeff = np.squeeze(self.dist_coeff)

        # 读取外参矩阵
        self.extrinsic_data = self.read_yaml('extrinsic_matrix.yaml')
        #if self.extrinsic_data is None:
        #    print("无法加载外参矩阵，程序终止。")
        #    return
        self.rotation_vector = self.check_and_reshape(self.extrinsic_data['rotation_vector'], (3,), "旋转向量")
        self.rotation_vector = np.squeeze(self.rotation_vector)
        self.translation_vector = self.check_and_reshape(self.extrinsic_data['translation_vector'], (3,), "平移向量")
        self.translation_vector = np.squeeze(self.translation_vector)

        # 假设世界坐标系中平面为Z=0平面
        self.world_plane_normal = np.array([0, 0, 1])


    def validate_vector(self, vec, name, length):
        if not isinstance(vec, np.ndarray) or vec.shape != (length,):
            raise ValueError(f"{name} 必须是形状为 ({length},) 的numpy数组")


    # 读取 YAML 文件
    def read_yaml(self, file_path):
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                data = yaml.safe_load(file)
            return data
        except FileNotFoundError:
            print(f"错误: 文件 {file_path} 未找到。")
            return None
        except yaml.YAMLError as e:
            print(f"错误: 解析 YAML 文件 {file_path} 时出错: {e}")
            return None


    # 检查并处理数据维度
    def check_and_reshape(self, data, target_shape, name):
        data = np.array(data)
        if data.shape != target_shape:
            data = data.reshape(target_shape)
            #print(f"警告: {name} 形状为 {data.shape}，已调整为 {target_shape}。")
        return data


    # 计算相机坐标系下的坐标
    def pixel_to_camera(self, pixel_coords, original_image_shape, camera_matrix, dist_coeff, rotation_vector, translation_vector,
                        world_plane_normal):
        # 验证输入参数
        #self.validate_vector(pixel_coords, "像素坐标", 2)
        #self.validate_vector(camera_matrix.reshape(-1), "相机内参矩阵", 9)
        #self.validate_vector(dist_coeff, "畸变系数", 5)
        #self.validate_vector(rotation_vector, "旋转向量", 3)
        #self.validate_vector(translation_vector, "平移向量", 3)
        #self.validate_vector(world_plane_normal, "世界平面法向量", 3)

        # 提取内参矩阵参数
        fx = camera_matrix[0, 0]
        fy = camera_matrix[1, 1]
        cx = camera_matrix[0, 2]
        cy = camera_matrix[1, 2]

        # 将世界坐标系下的平面法向量转换到相机坐标系
        rotation_matrix, _ = cv2.Rodrigues(rotation_vector)
        plane_normal_camera = np.dot(rotation_matrix, world_plane_normal.reshape(-1, 1)).flatten()

        # 计算平面在相机坐标系中的方程系数 (n·P + d = 0)
        # 假设世界坐标系中平面原点为(0,0,0)，转换到相机坐标系
        plane_point_camera = translation_vector  # 因为 R*0 + t = t
        d = -np.dot(plane_normal_camera, plane_point_camera)

        # 构造射线方向向量（归一化坐标）
        xn = (pixel_coords[0] - cx) / fx
        yn = (pixel_coords[1] - cy) / fy
        ray_dir = np.array([xn, yn, 1.0])

        # 计算射线与平面交点
        denominator = np.dot(plane_normal_camera, ray_dir)
        if abs(denominator) < 1e-6:
            raise ValueError("射线与平面几乎平行，无有效交点")

        t = -(np.dot(plane_normal_camera, [0, 0, 0]) + d) / denominator  # 相机原点在(0,0,0)
        camera_coords = t * ray_dir

        return camera_coords.reshape(1, 3)
    
    def get(self, x, y):
        pixel_coords = np.array([x, y], dtype=np.float32)
        camera_coords = self.pixel_to_camera(pixel_coords, None, self.camera_matrix, self.dist_coeff,
                                             self.rotation_vector, self.translation_vector, self.world_plane_normal)
        return camera_coords[0]





if __name__ == '__main__':
    loc = location()
    result = loc.get(123, 123)
    print(result)